from flask import Blueprint,request
from app.util.result import Result
from app.util.llm import llm_chunk,llm_embedding
from app.util.milvus import create_collection,insert_data,remove_collection,search_data
from langchain.text_splitter import CharacterTextSplitter
import requests
import io
import PyPDF2
import docx2txt


file_handle = Blueprint('file_handle', __name__)

"""
对文件进行嵌入
"""
@file_handle.route('/fileEmbedding',methods=['POST'])
def fileEmbedding():
    # 获取数据
    data = request.get_json()

    # 获取得到文件的url地址
    url = data['url']
    # 获取文件类型
    file_mime = data['fileMime']
    # 文件名
    file_name = data['fileName']
    resp = requests.get(url)
    
    #如果file_mime 不是 pdf txt docx md 的文件则返回错误
    if file_mime not in ['pdf','txt','docx','md']:
        return Result.error('文件类型错误').to_dict()
    # 文件内容
    file_content = resp.content
    bytes_content = io.BytesIO(file_content)
    if file_mime == 'pdf':
        # pdf 文件
        pdfReader = PyPDF2.PdfReader(bytes_content)
        file_content= ''
        for page in pdfReader.pages:
            file_content += page.extract_text()
    elif file_mime == 'docx':
        # word文档
        file_content = docx2txt.process(bytes_content)
        
    
    # 文件内容分块 chunk,通过换行符号进行分块
    splitter = CharacterTextSplitter(chunk_size=300, chunk_overlap=50)
    content= splitter.split_text(content)

    # 调用嵌入模型生成向量
    vectors= []
    for i in range(len(content)):
        vectors.append({
            "id": i+1,
            "vector": llm_embedding(content[i]),
            "text": content[i]
        })

    # 保存向量
    insert_data(file_name,vectors)

    # 获取文件类型
    return Result.success().to_dict()


@file_handle.route('/fileSearch',methods=['POST'])
def fileSearch():
    res_data = request.get_json()

    #搜索词
    query = res_data['query']
    #搜索文件名
    fileName = res_data['fileName']
    embedding= llm_embedding(query)
    # 搜索向量
    result =  search_data(fileName,embedding)
    return Result.success(result).to_dict()
    